package com.lg.lock;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;
import org.junit.Test;

import java.io.IOException;
import java.sql.Connection;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

/**
 * 基于zookeeper 的临时有序节点 实现的分布式锁，
 * 相比上一种，不需要循环监听
 * 此种方案是，多个线程抢占锁时，排序的线程监听比自己小1 的节点。这样一来可以说是实现了公平独占锁
 *
 * @author by Mr. Li 2020/12/10 15:48
 */
public class ZooKeeperDistributedLock implements Watcher {

    private ZooKeeper zk;
    private String locksRoot = "/locks";
    private Long productId;
    private String waitNode;
    private String lockNode;
    private CountDownLatch latch;
    private CountDownLatch connectedLatch = new CountDownLatch(1);
    private int sessionTimeout = 30000;
    private String connectString = "192.168.31.201:2181,192.168.31.201:2182,192.168.31.201:2183,192.168.31.201:2184,192.168.31.201:2185";


    public ZooKeeperDistributedLock(Long productId) {
        try {
            this.productId = productId;
            this.zk = new ZooKeeper(connectString, sessionTimeout, this);
            connectedLatch.await();
        } catch (IOException | InterruptedException e) {
            e.printStackTrace();
        }
    }

    public void acquireDistributedLock(Long path) {
        if (this.tryLock(path)) {
            return;
        } else {
            waitForLock(waitNode, sessionTimeout);
        }
    }


    /**
     * 尝试获取锁
     * <p>
     * 因为节点是顺序的，所以，需要获取所有的节点信息，并排列，然后根据顺讯进行获取锁。
     *
     * @return
     */
    private boolean tryLock(Long path) {
        try {
            lockNode = createPath(getPath(path));
            // 获取 locks 节点下所有的节点 进行排序
            List<String> locks = zk.getChildren(locksRoot, false);
            Collections.sort(locks);

            if (lockNode.equals(getPath(Long.valueOf(locks.get(0))))) {
                System.out.println("获取锁成功，path：" + lockNode);
                // 是当前的第一个节点，表示获取锁成功
                return true;
            }
            //如果不是最小的节点，找到比自己小1的节点
            int previousLockIndex = -1;
            for (int i = 0; i < locks.size(); i++) {
                if (lockNode.equals(locksRoot + "/" + locks.get(i))) {
                    previousLockIndex = i - 1;
                    break;
                }
            }
            this.waitNode = locks.get(previousLockIndex);
        } catch (InterruptedException | KeeperException e) {
            e.printStackTrace();
        }
        return false;
    }

    private boolean waitForLock(String waitNode, long waitTime) {
        // 判断等待的节点是否存在
        try {
            Stat stat = zk.exists(getPath(waitTime), true);
            if (stat != null) {
                this.latch = new CountDownLatch(1);
                this.latch.await(waitTime, TimeUnit.MILLISECONDS);
                this.latch = null;
            }
            return true;
        } catch (KeeperException | InterruptedException e) {
            return false;
        }

    }

    public void unlock() {
        try {
            // 删除/locks/10000000000节点
            // 删除/locks/10000000001节点
            System.out.println("unlock " + lockNode);
            zk.delete(lockNode, -1);
            lockNode = null;
            zk.close();
        } catch (InterruptedException | KeeperException e) {
            e.printStackTrace();
        }
    }
    public void unlock(Long path) {
        try {
            // 删除/locks/10000000000节点
            // 删除/locks/10000000001节点
            System.out.println("unlock " + path);
            zk.delete(getPath(path), -1);
            lockNode = null;
            zk.close();
        } catch (InterruptedException | KeeperException e) {
            e.printStackTrace();
        }
    }

    @Override
    public void process(WatchedEvent event) {
        if (Event.KeeperState.SyncConnected == event.getState()) {
            connectedLatch.countDown();
            System.out.println("zk连接状态成功,path:" + event.getPath());
            return;
        }
        if (this.latch != null) {
            this.latch.countDown();
        }
    }

    /**
     * 创建临时顺序节点
     *
     * @param path
     * @throws KeeperException
     * @throws InterruptedException
     */
    private String createPath(String path) throws KeeperException, InterruptedException {
        return zk.create(path, "".getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
    }

    private String getPath(Long productId) {
        return locksRoot + "/" + productId;
    }

    public static void main(String[] args) throws IOException {
        for (int i = 1; i < 10; i++) {
            ZooKeeperDistributedLock zooKeeperDistributedLock = new ZooKeeperDistributedLock(1L);
            zooKeeperDistributedLock.acquireDistributedLock(Long.valueOf(i));
        }
        System.in.read();
    }

    @Test
    public void test() throws KeeperException, InterruptedException {
        ZooKeeperDistributedLock zooKeeperDistributedLock = new ZooKeeperDistributedLock(1L);
        List<String> children = zk.getChildren(locksRoot, false);
        Collections.sort(children);

        zooKeeperDistributedLock.unlock(Long.valueOf(children.get(0)));
    }
}
